import com.android.build.api.transform.Context
import com.android.build.api.transform.DirectoryInput
import com.android.build.api.transform.Format
import com.android.build.api.transform.JarInput
import com.android.build.api.transform.QualifiedContent
import com.android.build.api.transform.SecondaryInput
import com.android.build.api.transform.Transform
import com.android.build.api.transform.TransformInput
import com.android.build.api.transform.TransformInvocation
import com.android.build.api.transform.TransformOutputProvider
import org.gradle.api.Project
import org.gradle.api.logging.Logger
import org.gradle.api.logging.Logging
import org.objectweb.asm.ClassReader
import org.objectweb.asm.ClassWriter
import org.objectweb.asm.tree.ClassNode
import utils.FileUtils
import java.io.File
import java.io.FileOutputStream
import java.io.InputStream
import java.util.LinkedList
import java.util.concurrent.CountDownLatch
import java.util.concurrent.ExecutorService
import java.util.concurrent.Executors
import java.util.jar.JarEntry
import java.util.jar.JarFile
import java.util.jar.JarOutputStream


/**
 * Author: liaohailong
 * Date: 2024/3/25
 * Time: 21:11
 * Description: 第一个自定义的 Transformer
 **/
class CustomTransformer(private val project: Project) : Transform() {
    private val logger: Logger by lazy { Logging.getLogger("CustomTransformer") }
    private fun log(msg: String) {
        logger.quiet("${name}: $msg")
    }

    override fun getName() = "lookLook"

    override fun getInputTypes() = hashSetOf(QualifiedContent.DefaultContentType.CLASSES) // 只关心字节码

    override fun getScopes() = hashSetOf(
        QualifiedContent.Scope.PROJECT, // 本工程
        QualifiedContent.Scope.SUB_PROJECTS, // 本工程依赖的子工程（模块）
        QualifiedContent.Scope.EXTERNAL_LIBRARIES // jar包 aar包
    )

    override fun isIncremental() = true

    override fun transform(transformInvocation: TransformInvocation) {
        log("Start -------> transform")
        log(">>> gradle version: ${project.gradle.gradleVersion}")
        log(">>> isIncremental: ${transformInvocation.isIncremental}")


        val costMs = measureTransformMilliseconds(transformInvocation) { context: Context,
                                                                         inputs: Collection<TransformInput>,
                                                                         referencedInputs: Collection<TransformInput>,
                                                                         secondaryInputs: Collection<SecondaryInput>,
                                                                         outputProvider: TransformOutputProvider,
                                                                         incremental: Boolean ->

            // 非增量编译的情况
            if (!incremental) outputProvider.deleteAll()

            val outputJarList = LinkedList<JarInfo>() // jar包中的内容
            val outputDirList = LinkedList<DirItemInfo>() // 源码编译出来的class

            // 处理输入数据
            inputs.forEach {

                val jarInputs = it.jarInputs // 项目中使用到的jar包
                val directoryInputs = it.directoryInputs // 源码编译出来的class所在的目录

                val startMs = System.currentTimeMillis()
                // 字节码扫描：编译所有上游传入的内容，进行分类处理
                val count = process1(
                    jarInputs,
                    directoryInputs,
                    outputProvider,
                    outputJarList,
                    outputDirList
                )

                val costMs = System.currentTimeMillis() - startMs
                log("last process cost ${costMs}ms, $count file scanned")
            }

            // 字节码修改与输出
            process2(outputJarList, outputDirList)
        }

        log("End -------> transform cost: ${costMs}ms")
    }

    private fun measureTransformMilliseconds(
        transformInvocation: TransformInvocation,
        block: (
            context: Context,
            inputs: Collection<TransformInput>,
            referencedInputs: Collection<TransformInput>,
            secondaryInputs: Collection<SecondaryInput>,
            outputProvider: TransformOutputProvider,
            incremental: Boolean
        ) -> Unit
    ): Long {
        // gradle运行的上下文，可以拿到日志等级，名称路径等信息
        val context = transformInvocation.context
        // 本次是否为增量编译
        val incremental = transformInvocation.isIncremental

        // 配合增量的输入
        val secondaryInputs = transformInvocation.secondaryInputs

        // 引用输入，该输入处理后，不需要输出
        val referencedInputs = transformInvocation.referencedInputs

        // 普通输入，该输入处理完毕之后需要输出到指定位置上
        val inputs = transformInvocation.inputs

        // 用于提供输入的位置，这里提供成输出的内容，会成为下一个transformer的输入
        val outputProvider = transformInvocation.outputProvider

        val startMs = System.currentTimeMillis()
        block(context, inputs, referencedInputs, secondaryInputs, outputProvider, incremental)
        return System.currentTimeMillis() - startMs
    }

    /**
     * 预处理
     */
    private fun process1(
        jis: Collection<JarInput>, // jarInputs
        dis: Collection<DirectoryInput>, // directoryInputs
        output: TransformOutputProvider, // 输出控制
        outputJarList: LinkedList<JarInfo>, // jar包中的内容
        outputDirList: LinkedList<DirItemInfo>, // 源码编译产物，一般是class
    ): Long {
        // 扫描出来的文件计数
        var count = 0L

        // 把输入源的jar包，配置到输出上
        jis.forEach { input ->
            val src = input.file
            // 获取输出jar的位置，实际上是一个file，后续要把jar包内容拷贝过去
            val dest = output.getContentLocation(
                input.name,
                input.contentTypes,
                input.scopes,
                Format.JAR
            )

            // jar包内容
            JarFile(src).use { jf ->
                val entries = jf.entries()
                val list = LinkedList<JarItemInfo>()
                // 这里可以做一些解析操作，把字节码的文件流，交给ASM处理
                for (entry in entries) {
                    if (entry.isUseful()) {
                        val classInfo = jf.getInputStream(entry).visit()
                        list.add(JarItemInfo(entry, classInfo)) // 可能需要字节码修改
//                        log("process jar class: ${entry.realName}")
                    } else {
                        list.add(JarItemInfo(entry)) // 无需字节码修改
                    }
                    count++
                }
                outputJarList.add(JarInfo(src, dest, list))
            }
        }

        dis.forEach { input ->
            val src = input.file
            // 获取源码编译产物目录（class和其他配置文件存放的目录）
            val dest = output.getContentLocation(
                input.name,
                input.contentTypes,
                input.scopes,
                Format.DIRECTORY
            )
            // 源码编译出来的文件全部放进去，不影响下一个transformer
            FileUtils.copyDirectory(src, dest)
            dest.walk().forEach {
                // 只对某些关心的文件做处理，这里暂时只关心class文件
                if (it.isUseful()) {
                    val classInfo = it.inputStream().visit()
                    outputDirList.add(DirItemInfo(it, classInfo)) // 可能需要字节码修改
//                    log("process dir file: ${it.name}")
                } else {
                    outputDirList.add(DirItemInfo(it)) // 无需字节码修改
                }
                count++
            }
        }

        return count
    }

    /**
     * 正在处理
     */
    private fun process2(
        outputJarList: LinkedList<JarInfo>, // jar包中的内容
        outputDirList: LinkedList<DirItemInfo>, // 源码编译产物，一般是class
    ) {
        // TODO: 执行钩子函数，进行字节码处理

        // 多线程处理
        val executors = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors())
        var throwable: Throwable? = null
        val dirLatch = CountDownLatch(outputDirList.size)

        // 源码产物
        outputDirList.forEach { item ->
            executors.exec(dirLatch, onError = { throwable = it }) {
                val ci = item.classInfo
                if (ci != null) {
                    // 把修改好的字节码同步写入到输出中去
                    val cw = ClassWriter(ci.classReader, 0)
                    ci.classNode.accept(cw)
                    FileOutputStream(item.dest).use { it.write(cw.toByteArray()) }
                }
            }
        }


        dirLatch.await()
        throwable?.also { executors.shutdownNow();throw it }


        // jar包产物
        val jarLatch = CountDownLatch(outputJarList.size)
        outputJarList.forEach { jarInfo ->
            executors.exec(jarLatch, onError = { throwable = it }) {
                // 把jar包中的内容写到输出中去
                JarOutputStream(jarInfo.dest.outputStream()).use { jos ->
                    jarInfo.list.forEach { entity ->
                        val ci = entity.classInfo
                        if (ci == null) {
                            // 无字节码修改的内容
                            JarFile(jarInfo.src).use { jf ->
                                jos.putNextEntry(JarEntry(entity.data.name))
                                jos.write(jf.getInputStream(entity.data).readBytes())
                                jos.closeEntry()
                            }
                        } else {
                            // 可能出现字节码修改的内容
                            val cw = ClassWriter(ci.classReader, 0)
                            ci.classNode.accept(cw)
                            jos.putNextEntry(JarEntry(entity.data.name))
                            jos.write(cw.toByteArray())
                            jos.closeEntry()
                        }
                    }
                }

            }
        }

        jarLatch.await()
        executors.shutdownNow()
        throwable?.also { executors.shutdownNow();throw it }

    }


    /**
     * @return 返回true表示，该文件是一个class文件（Android自动生成文件除外）
     */
    private fun JarEntry.isUseful(): Boolean {
        return name.endsWith("class") && !checkAndroidRFile(name)
    }

    /**
     * @return 返回true表示，该文件是一个class文件
     */
    private fun File.isUseful(): Boolean = isFile && extension == "class"

    private fun InputStream.visit() = use {
        val cr = ClassReader(it)
        val cn = ClassNode()
        cr.accept(cn, 0)

        // TODO: 交给钩子函数，执行字节码逻辑初始化

        ClassInfo(cr, cn)
    }

    /**
     * @param src 源文件jar
     * @param dest 输出文件jar
     * @param list jar中的子节点内容，可能是class，可能是 META-INF/
     */
    private class JarInfo(
        val src: File,
        val dest: File,
        val list: LinkedList<JarItemInfo>,
    )

    /**
     * @param data Jar包中的子节点内容
     * @param classInfo 准备使用ASM，做字节码修改
     */
    private class JarItemInfo(
        val data: JarEntry,
        val classInfo: ClassInfo? = null
    )

    /**
     * @param dest 输出的源码编译产物，一般是class
     * @param classInfo 准备使用ASM，做字节码修改
     */
    private class DirItemInfo(
        val dest: File,
        val classInfo: ClassInfo? = null
    )

    /**
     * ASM字节码修改相关内容
     */
    private class ClassInfo(
        val classReader: ClassReader,
        val classNode: ClassNode
    )


    private fun ExecutorService.exec(
        counter: CountDownLatch,
        onError: (throwable: Throwable) -> Unit,
        runnable: Runnable
    ) {
        execute {
            try {
                runnable.run()
            } catch (ex: Exception) {
                onError(ex)
            } finally {
                counter.countDown()
            }
        }
    }

}